package net.sourceforge.jsocks;

import net.sourceforge.jsocks.server.ServerAuthenticator;

import java.io.EOFException;
import java.io.IOException;
import java.io.InputStream;
import java.io.InterruptedIOException;
import java.io.OutputStream;
import java.io.PrintStream;
import java.io.PushbackInputStream;
import java.net.ConnectException;
import java.net.InetAddress;
import java.net.NoRouteToHostException;
import java.net.ServerSocket;
import java.net.Socket;

/**
 * SOCKS4 and SOCKS5 proxy, handles both protocols simultaniously. Implements
 * all SOCKS commands, including UDP relaying.
 * <p>
 * In order to use it you will need to implement ServerAuthenticator interface.
 * There is an implementation of this interface which does no authentication
 * ServerAuthenticatorNone, but it is very dangerous to use, as it will give
 * access to your local network to anybody in the world. One should never use
 * this authentication scheme unless one have pretty good reason to do so. There
 * is a couple of other authentication schemes in socks.server package.
 *
 * @see socks.server.ServerAuthenticator
 */
public class ProxyServer implements Runnable {

    static final int START_MODE = 0;
    static final int ACCEPT_MODE = 1;
    static final int PIPE_MODE = 2;
    static final int ABORT_MODE = 3;
    static final int BUF_SIZE = 8192;
    static final String command_names[] = {"CONNECT", "BIND", "UDP_ASSOCIATE"};
    protected static int iddleTimeout = 180000; // 3 minutes
    static int acceptTimeout = 180000; // 3 minutes
    static PrintStream log = null;
    static Proxy proxy;
    ServerAuthenticator auth;
    ProxyMessage msg = null;
    Socket sock = null, remote_sock = null;
    ServerSocket ss = null;
    UDPRelayServer relayServer = null;
    InputStream in, remote_in;
    OutputStream out, remote_out;
    int mode;
    Thread pipe_thread1, pipe_thread2;

    // Public Constructors
    // ///////////////////
    long lastReadTime;

    // Other constructors
    // //////////////////

    /**
     * Creates a proxy server with given Authentication scheme.
     *
     * @param auth Authentication scheme to be used.
     */
    public ProxyServer(ServerAuthenticator auth) {
        this.auth = auth;
    }

    // Public methods
    // ///////////////

    protected ProxyServer(ServerAuthenticator auth, Socket s) {
        this.auth = auth;
        this.sock = s;
        mode = START_MODE;
    }

    static final String command2String(int cmd) {
        if (cmd > 0 && cmd < 4)
            return command_names[cmd - 1];
        else
            return "Unknown Command " + cmd;
    }

    /**
     * Get proxy.
     *
     * @return Proxy wich is used to handle user requests.
     */
    public static Proxy getProxy() {
        return proxy;
    }

    /**
     * Set proxy.
     * <p>
     * Allows Proxy chaining so that one Proxy server is connected to another
     * and so on. If proxy supports SOCKSv4, then only some SOCKSv5 requests can
     * be handled, UDP would not work, however CONNECT and BIND will be
     * translated.
     *
     * @param p Proxy which should be used to handle user requests.
     */
    public static void setProxy(Proxy p) {
        proxy = p;
        UDPRelayServer.proxy = proxy;
    }

    static final void log(ProxyMessage msg) {
        log("Request version:" + msg.version + "\tCommand: "
                + command2String(msg.command));
        log("IP:" + msg.ip + "\tPort:" + msg.port
                + (msg.version == 4 ? "\tUser:" + msg.user : ""));
    }

    static final void log(String s) {
        if (log != null) {
            log.println(s);
            log.flush();
        }
    }

    /**
     * Sets the timeout for BIND command, how long the server should wait for
     * the incoming connection.<br>
     * Zero timeout implies infinity.<br>
     * Default timeout is 3 minutes.
     */
    public static void setAcceptTimeout(int timeout) {
        acceptTimeout = timeout;
    }

    /**
     * Sets the size of the datagrams used in the UDPRelayServer.<br>
     * Default size is 64K, a bit more than maximum possible size of the
     * datagram.
     */
    public static void setDatagramSize(int size) {
        UDPRelayServer.setDatagramSize(size);
    }

    /**
     * Sets the timeout for connections, how long shoud server wait for data to
     * arrive before dropping the connection.<br>
     * Zero timeout implies infinity.<br>
     * Default timeout is 3 minutes.
     */
    public static void setIddleTimeout(int timeout) {
        iddleTimeout = timeout;
    }

    /**
     * Set the logging stream. Specifying null disables logging.
     */
    public static void setLog(OutputStream out) {
        if (out == null) {
            log = null;
        } else {
            log = new PrintStream(out, true);
        }

        UDPRelayServer.log = log;
    }

    /**
     * Sets the timeout for UDPRelay server.<br>
     * Zero timeout implies infinity.<br>
     * Default timeout is 3 minutes.
     */
    public static void setUDPTimeout(int timeout) {
        UDPRelayServer.setTimeout(timeout);
    }

    private synchronized void abort() {
        if (mode == ABORT_MODE)
            return;
        mode = ABORT_MODE;
        try {
            log("Aborting operation");
            if (remote_sock != null)
                remote_sock.close();
            if (sock != null)
                sock.close();
            if (relayServer != null)
                relayServer.stop();
            if (ss != null)
                ss.close();
            if (pipe_thread1 != null)
                pipe_thread1.interrupt();
            if (pipe_thread2 != null)
                pipe_thread2.interrupt();
        } catch (IOException ioe) {
        }
    }

    private void doAccept() throws IOException {
        Socket s;
        long startTime = System.currentTimeMillis();

        while (true) {
            s = ss.accept();
            if (s.getInetAddress().equals(msg.ip)) {
                // got the connection from the right host
                // Close listenning socket.
                ss.close();
                break;
            } else if (ss instanceof SocksServerSocket) {
                // We can't accept more then one connection
                s.close();
                ss.close();
                throw new SocksException(Proxy.SOCKS_FAILURE);
            } else {
                if (acceptTimeout != 0) { // If timeout is not infinit
                    int newTimeout = acceptTimeout
                            - (int) (System.currentTimeMillis() - startTime);
                    if (newTimeout <= 0)
                        throw new InterruptedIOException("In doAccept()");
                    ss.setSoTimeout(newTimeout);
                }
                s.close(); // Drop all connections from other hosts
            }
        }

        // Accepted connection
        remote_sock = s;
        remote_in = s.getInputStream();
        remote_out = s.getOutputStream();

        // Set timeout
        remote_sock.setSoTimeout(iddleTimeout);

        log("Accepted from " + s.getInetAddress() + ":" + s.getPort());

        ProxyMessage response;

        if (msg.version == 5)
            response = new Socks5Message(Proxy.SOCKS_SUCCESS,
                    s.getInetAddress(), s.getPort());
        else
            response = new Socks4Message(Socks4Message.REPLY_OK,
                    s.getInetAddress(), s.getPort());
        response.write(out);
    }

    private void handleException(IOException ioe) {
        // If we couldn't read the request, return;
        if (msg == null)
            return;
        // If have been aborted by other thread
        if (mode == ABORT_MODE)
            return;
        // If the request was successfully completed, but exception happened
        // later
        if (mode == PIPE_MODE)
            return;

        int error_code = Proxy.SOCKS_FAILURE;

        if (ioe instanceof SocksException)
            error_code = ((SocksException) ioe).errCode;
        else if (ioe instanceof NoRouteToHostException)
            error_code = Proxy.SOCKS_HOST_UNREACHABLE;
        else if (ioe instanceof ConnectException)
            error_code = Proxy.SOCKS_CONNECTION_REFUSED;
        else if (ioe instanceof InterruptedIOException)
            error_code = Proxy.SOCKS_TTL_EXPIRE;

        if (error_code > Proxy.SOCKS_ADDR_NOT_SUPPORTED || error_code < 0) {
            error_code = Proxy.SOCKS_FAILURE;
        }

        sendErrorMessage(error_code);
    }

    protected void handleRequest(ProxyMessage msg) throws IOException {
        if (!auth.checkRequest(msg))
            throw new SocksException(Proxy.SOCKS_FAILURE);

        if (msg.ip == null) {
            if (msg instanceof Socks5Message) {
                msg.ip = InetAddress.getByName(msg.host);
            } else
                throw new SocksException(Proxy.SOCKS_FAILURE);
        }
        log(msg);

        switch (msg.command) {
            case Proxy.SOCKS_CMD_CONNECT:
                onConnect(msg);
                break;
            case Proxy.SOCKS_CMD_BIND:
                onBind(msg);
                break;
            case Proxy.SOCKS_CMD_UDP_ASSOCIATE:
                onUDP(msg);
                break;
            default:
                throw new SocksException(Proxy.SOCKS_CMD_NOT_SUPPORTED);
        }
    }

    private void onBind(ProxyMessage msg) throws IOException {
        ProxyMessage response = null;

        if (proxy == null)
            ss = new ServerSocket(0);
        else
            ss = new SocksServerSocket(proxy, msg.ip, msg.port);

        ss.setSoTimeout(acceptTimeout);

        log("Trying accept on " + ss.getInetAddress() + ":" + ss.getLocalPort());

        if (msg.version == 5)
            response = new Socks5Message(Proxy.SOCKS_SUCCESS,
                    ss.getInetAddress(), ss.getLocalPort());
        else
            response = new Socks4Message(Socks4Message.REPLY_OK,
                    ss.getInetAddress(), ss.getLocalPort());
        response.write(out);

        mode = ACCEPT_MODE;

        pipe_thread1 = Thread.currentThread();
        pipe_thread2 = new Thread(this);
        pipe_thread2.start();

        // Make timeout infinit.
        sock.setSoTimeout(0);
        int eof = 0;

        try {
            while ((eof = in.read()) >= 0) {
                if (mode != ACCEPT_MODE) {
                    if (mode != PIPE_MODE)
                        return;// Accept failed

                    remote_out.write(eof);
                    break;
                }
            }
        } catch (EOFException eofe) {
            // System.out.println("EOF exception");
            return;// Connection closed while we were trying to accept.
        } catch (InterruptedIOException iioe) {
            // Accept thread interrupted us.
            // System.out.println("Interrupted");
            if (mode != PIPE_MODE)
                return;// If accept thread was not successfull return.
        } finally {
            // System.out.println("Finnaly!");
        }

        if (eof < 0)// Connection closed while we were trying to accept;
            return;

        // Do not restore timeout, instead timeout is set on the
        // remote socket. It does not make any difference.

        pipe(in, remote_out);
    }

    private void onConnect(ProxyMessage msg) throws IOException {
        Socket s;
        ProxyMessage response = null;

        s = new Socket(msg.ip, msg.port);

        log("Connected to " + s.getInetAddress() + ":" + s.getPort());

        if (msg instanceof Socks5Message) {
            response = new Socks5Message(Proxy.SOCKS_SUCCESS,
                    s.getLocalAddress(), s.getLocalPort());
        } else {
            response = new Socks4Message(Socks4Message.REPLY_OK,
                    s.getLocalAddress(), s.getLocalPort());

        }
        response.write(out);
        startPipe(s);
    }

    // Private methods
    // ////////////////

    private void onUDP(ProxyMessage msg) throws IOException {
        if (msg.ip.getHostAddress().equals("0.0.0.0"))
            msg.ip = sock.getInetAddress();
        log("Creating UDP relay server for " + msg.ip + ":" + msg.port);
        relayServer = new UDPRelayServer(msg.ip, msg.port,
                Thread.currentThread(), sock, auth);

        ProxyMessage response;

        response = new Socks5Message(Proxy.SOCKS_SUCCESS, relayServer.relayIP,
                relayServer.relayPort);

        response.write(out);

        relayServer.start();

        // Make timeout infinit.
        sock.setSoTimeout(0);
        try {
            while (in.read() >= 0)
                /* do nothing */ ;
        } catch (EOFException eofe) {
        }
    }

    private void pipe(InputStream in, OutputStream out) throws IOException {
        lastReadTime = System.currentTimeMillis();
        byte[] buf = new byte[BUF_SIZE];
        int len = 0;
        while (len >= 0) {
            try {
                if (len != 0) {
                    out.write(buf, 0, len);
                    out.flush();
                }
                len = in.read(buf);
                lastReadTime = System.currentTimeMillis();
            } catch (InterruptedIOException iioe) {
                if (iddleTimeout == 0)
                    return;// Other thread interrupted us.
                long timeSinceRead = System.currentTimeMillis() - lastReadTime;
                if (timeSinceRead >= iddleTimeout - 1000) // -1s for adjustment.
                    return;
                len = 0;

            }
        }
    }

    protected ProxyMessage readMsg(InputStream in) throws IOException {
        PushbackInputStream push_in;
        if (in instanceof PushbackInputStream)
            push_in = (PushbackInputStream) in;
        else
            push_in = new PushbackInputStream(in);

        int version = push_in.read();
        push_in.unread(version);

        ProxyMessage msg;

        if (version == 5) {
            msg = new Socks5Message(push_in, false);
        } else if (version == 4) {
            msg = new Socks4Message(push_in, false);
        } else {
            throw new SocksException(Proxy.SOCKS_FAILURE);
        }
        return msg;
    }

    // Runnable interface
    // //////////////////
    @Override
    public void run() {
        switch (mode) {
            case START_MODE:
                try {
                    startSession();
                } catch (IOException ioe) {
                    handleException(ioe);
                    // ioe.printStackTrace();
                } finally {
                    abort();
                    if (auth != null)
                        auth.endSession();
                    log("Main thread(client->remote)stopped.");
                }
                break;
            case ACCEPT_MODE:
                try {
                    doAccept();
                    mode = PIPE_MODE;
                    pipe_thread1.interrupt(); // Tell other thread that connection
                    // have
                    // been accepted.
                    pipe(remote_in, out);
                } catch (IOException ioe) {
                    // log("Accept exception:"+ioe);
                    handleException(ioe);
                } finally {
                    abort();
                    log("Accept thread(remote->client) stopped");
                }
                break;
            case PIPE_MODE:
                try {
                    pipe(remote_in, out);
                } catch (IOException ioe) {
                } finally {
                    abort();
                    log("Support thread(remote->client) stopped");
                }
                break;
            case ABORT_MODE:
                break;
            default:
                log("Unexpected MODE " + mode);
        }
    }

    private void sendErrorMessage(int error_code) {
        ProxyMessage err_msg;
        if (msg instanceof Socks4Message)
            err_msg = new Socks4Message(Socks4Message.REPLY_REJECTED);
        else
            err_msg = new Socks5Message(error_code);
        try {
            err_msg.write(out);
        } catch (IOException ioe) {
        }
    }

    /**
     * Start the Proxy server at given port.<br>
     * This methods blocks.
     */
    public void start(int port) {
        start(port, 5, null);
    }

    /**
     * Create a server with the specified port, listen backlog, and local IP
     * address to bind to. The localIP argument can be used on a multi-homed
     * host for a ServerSocket that will only accept connect requests to one of
     * its addresses. If localIP is null, it will default accepting connections
     * on any/all local addresses. The port must be between 0 and 65535,
     * inclusive. <br>
     * This methods blocks.
     */
    public void start(int port, int backlog, InetAddress localIP) {
        try {
            ss = new ServerSocket(port, backlog, localIP);
            log("Starting SOCKS Proxy on:"
                    + ss.getInetAddress().getHostAddress() + ":"
                    + ss.getLocalPort());
            while (true) {
                Socket s = ss.accept();
                log("Accepted from:" + s.getInetAddress().getHostName() + ":"
                        + s.getPort());
                ProxyServer ps = new ProxyServer(auth, s);
                (new Thread(ps)).start();
            }
        } catch (IOException ioe) {
            ioe.printStackTrace();
        } finally {
        }
    }

    private void startPipe(Socket s) {
        mode = PIPE_MODE;
        remote_sock = s;
        try {
            remote_in = s.getInputStream();
            remote_out = s.getOutputStream();
            pipe_thread1 = Thread.currentThread();
            pipe_thread2 = new Thread(this);
            pipe_thread2.start();
            pipe(in, remote_out);
        } catch (IOException ioe) {
        }
    }

    // Private methods
    // ///////////////
    private void startSession() throws IOException {
        sock.setSoTimeout(iddleTimeout);

        try {
            auth = auth.startSession(sock);
        } catch (IOException ioe) {
            log("Auth throwed exception:" + ioe);
            auth = null;
            return;
        }

        if (auth == null) { // Authentication failed
            log("Authentication failed");
            return;
        }

        in = auth.getInputStream();
        out = auth.getOutputStream();

        msg = readMsg(in);
        handleRequest(msg);
    }

    /**
     * Stop server operation.It would be wise to interrupt thread running the
     * server afterwards.
     */
    public void stop() {
        try {
            if (ss != null)
                ss.close();
        } catch (IOException ioe) {
        }
    }
}
